"""
1. 设计一个提取backbone特征的hook

2. 设计一个context manager
    class ContextManager(...):
        self._hooks = ...

        .........

        @contextmanager
        def collect_context(...)
            收集上下文信息：backbone的特征以及传递的task id
            存储上下文信息

        .........

        @contextmanager
        def task_context(task_ids: Optional[List[int]] = None, clear:bool=True):
            传递上下文信息
            是否清空存储的上下文信息

        ...

"""

"""
Context manager for collecting and managing backbone features and task information.

This module provides a unified way to collect backbone features and manage task contexts
during forward passes, enabling efficient feature sharing and task-aware processing.
"""

import weakref
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn

from continuallearning.interfaces import HookInterface, HookType
from continuallearning.models.pefts.common import ModuleDetector
from continuallearning.models.pefts.common.config import (
    ContextManagerConfig,
)
from continuallearning.models.pefts.common.context import (
    add_backbone_feature,
    get_global_context,
)
from continuallearning.models.pefts.hook_managers._base_hook_manager import (
    AbstractHookManager,
)
from continuallearning.models.pefts.hooks._base_hook import HookBehaviour
from continuallearning.registry import HOOK_MANAGER_REGISTRY
from continuallearning.utils.logging import get_logger
from continuallearning.utils.types import EXTRACT_TENSOR_FROM_INPUT

logger = get_logger(__name__)


class BackboneFeatureHook(HookBehaviour, HookInterface):
    """
    Hook for extracting backbone features during forward pass.

    This hook captures intermediate features from backbone layers
    and stores them in the context manager for later use.
    """

    _hook_type: HookType = HookType.FORWARD_POST

    def __init__(
        self,
        module: nn.Module,
        layer_name: str,
        extract_input: bool = False,
        extract_output: bool = True,
    ):
        """
        Initialize the backbone feature hook.

        Args:
            module: The module to hook
            layer_name: Name identifier for this layer
            context_manager: Reference to the context manager
            extract_input: Whether to extract input features
            extract_output: Whether to extract output features
        """
        self._module_ref = weakref.ref(module)
        self.layer_name = layer_name
        self.extract_input = extract_input
        self.extract_output = extract_output

        # Use global context manager implicitly
        self._global_context = get_global_context()

    def _post_forward_hook(
        self, module: nn.Module, inputs: Tuple[torch.Tensor, ...], output: torch.Tensor
    ) -> torch.Tensor:
        """Forward hook using global context."""
        # Only collect if global context is collecting
        if not self._global_context.is_collecting:
            return output

        try:
            if self.extract_input and inputs:
                input_tensor = EXTRACT_TENSOR_FROM_INPUT(inputs)
                add_backbone_feature(f"{self.layer_name}_input", input_tensor)

            if self.extract_output:
                output_tensor = (
                    output if isinstance(output, torch.Tensor) else output[0]
                )
                add_backbone_feature(self.layer_name, output_tensor)

        except Exception as e:
            logger.error(f"Error in backbone feature hook {self.layer_name}: {e}")

        return output

    def get_parameters(self) -> List[nn.Parameter]:
        return []

    def get_trainable_parameters(self) -> List[nn.Parameter]:
        return []


@HOOK_MANAGER_REGISTRY.register("context")
class ContextHookManager(AbstractHookManager):
    """
    Context manager for collecting backbone features and managing task contexts.

    This class provides context managers for:
    1. Collecting backbone features during forward passes
    2. Managing task-specific contexts with automatic cleanup
    """

    def __init__(
        self,
        model: nn.Module,
        config: Union[Dict[str, Any], ContextManagerConfig] = {
            "discovery": {
                "path_patterns": [
                    # 匹配所有编码器层 (encoder.layer.0, encoder.layer.1, ...)
                    r"^encoder\.layer\.\d+$",
                ],
                "ignore_module_types": [
                    "Dropout",
                    "Identity",
                ],
            },
            "extract_output": True,
            "extract_input": False,
        },
    ):
        """Initialize the context manager."""
        self._model_ref = weakref.ref(model)
        self.config = self._process_config(config)
        self.module_detector = ModuleDetector(self.model, self.config.discovery)
        # Get global context reference
        self._global_context = get_global_context()

        # FIXME hook manager 继承关系和基类设计需要修改
        self._hooks: Dict[str, BackboneFeatureHook] = {}  # type: ignore

    def _process_config(
        self, config: Optional[Union[Dict[str, Any], ContextManagerConfig]]
    ) -> ContextManagerConfig:
        """Process and validate configuration."""
        if isinstance(config, ContextManagerConfig):
            config.validate()
            return config
        elif isinstance(config, dict):
            try:
                context_config = ContextManagerConfig(**config)
                context_config.validate()
                return context_config
            except Exception as e:
                raise ValueError(f"Invalid configuration dictionary: {e}") from e
        else:
            raise TypeError(
                f"config must be dict, or ContextManagerConfig, "
                f"got {type(config).__name__}"
            )

    @property
    def model(self) -> nn.Module:
        """
        Get the model instance.

        Returns:
            nn.Module: The model instance if still alive, None otherwise
        """
        if self._model_ref is None:
            raise RuntimeError("Model reference has not been set")

        model = self._model_ref()
        if model is None:
            raise RuntimeError("Model has been garbage collected")
        return model

    def add_hooks(self) -> None:
        """
        Automatically discover and attach backbone feature hooks based on configuration.

        This method performs the following steps:
        1. Finds target modules using the discovery configuration
        2. Creates BackboneFeatureHook for each target module
        3. Attaches each hook and tracks it
        """
        if self._is_attached:
            logger.warning("Hooks already attached, skipping")
            return

        # Find target modules based on configuration
        logger.info(
            f"Searching for backbone modules in model {type(self.model).__name__}"
        )
        target_modules = self.module_detector.find_target_modules()

        if not target_modules:
            logger.warning("No target modules found for backbone feature extraction")
            return

        logger.info(
            f"Found {len(target_modules)} backbone modules, attaching feature hooks..."
        )

        success_count = 0
        for idx, module_dict in enumerate(target_modules):
            path, module = module_dict["path"], module_dict["module"]
            module_type = type(module).__name__

            # Skip if this module already has a hook
            if path in self._hooks:
                logger.debug(
                    f"Module '{path}' [{module_type}] already has feature hook, skipping"
                )
                continue

            logger.debug(
                f"[{idx + 1}/{len(target_modules)}] Attaching feature hook to '{path}' [{module_type}]"
            )

            try:
                hook = BackboneFeatureHook(
                    module=module,
                    layer_name=path,
                    extract_input=self.config.extract_input,
                    extract_output=self.config.extract_output,
                )
                hook.attach()
                self._hooks[path] = hook
                success_count += 1

                logger.debug(f"✓ Successfully attached feature hook to '{path}'")

            except Exception as e:
                logger.error(
                    f"✗ Failed to attach feature hook to '{path}' [{module_type}]: {e}"
                )

        if success_count > 0:
            logger.info(
                f"Successfully attached {success_count}/{len(target_modules)} backbone feature hooks"
            )
            self._is_attached = True
        else:
            logger.warning(
                "No feature hooks were successfully attached, check configuration or model structure"
            )

    def remove_hooks(self) -> None:
        """
        Detach all backbone feature hooks.

        This method safely detaches all hooks managed by this context manager.
        """
        if not self._hooks:
            logger.debug("No backbone feature hooks to detach")
            return

        success_count = 0
        total_hooks = len(self._hooks)
        logger.info(f"Detaching all backbone feature hooks, total: {total_hooks}")

        for path, hook in list(self._hooks.items()):
            try:
                hook.detach()
                success_count += 1
                logger.debug(f"✓ Detached feature hook from '{path}'")
            except Exception as e:
                logger.error(f"✗ Failed to detach feature hook from '{path}': {e}")

        self._hooks.clear()
        self._is_attached = False
        logger.info(
            f"Detached backbone feature hooks, {success_count}/{total_hooks} successful"
        )

    def attach_target_layer(
        self,
        module: nn.Module,
        layer_name: str,
        extract_input: bool = False,
        extract_output: bool = True,
    ) -> "ContextHookManager":
        """
        Register a backbone layer for feature extraction.

        Args:
            module: The module to hook
            layer_name: Name identifier for this layer
            extract_input: Whether to extract input features
            extract_output: Whether to extract output features

        Returns:
            Self for method chaining
        """
        if layer_name in self._hooks:
            logger.warning(f"Layer {layer_name} already registered")
            return self

        hook = BackboneFeatureHook(
            module=module,
            layer_name=layer_name,
            extract_input=extract_input,
            extract_output=extract_output,
        )
        hook.attach()
        self._hooks[layer_name] = hook

        logger.debug(f"Registered backbone layer: {layer_name}")
        return self

    def detach_target_layer(self, layer_name: str) -> "ContextHookManager":
        """
        Unregister a backbone layer.

        Args:
            layer_name: Name of the layer to unregister

        Returns:
            Self for method chaining
        """
        if layer_name in self._hooks:
            self._hooks[layer_name].detach()
            del self._hooks[layer_name]
            logger.debug(f"Unregistered backbone layer: {layer_name}")
        else:
            logger.warning(f"Layer {layer_name} not found in registered hooks")
        return self

    def __repr__(self) -> str:
        """String representation of the context hook manager."""
        model_name = type(self.model).__name__ if self._model_ref else "None"
        return (
            f"ContextHookManager(model={model_name}, "
            f"hooks={len(self._hooks)}, attached={self._is_attached})"
        )
